# -*- coding: utf-8 -*-
# ㊎ Author : xmz
# ㊍ Date:  : 2025/1/4 下午3:07
# ㊌ File   : Snowflake.py
# ㊋ IDE    : PyCharm
# ㊏ REMARKS:
import os
import sys
import threading
import time

path = os.path.abspath(__file__)
sys.path.extend([path, os.path.dirname(path), os.path.dirname(os.path.dirname(path))])


class SnowflakeIDGenerator:

    def __init__(self, datacenter_id, worker_id):
        # 调整位数
        self.worker_id_bits = 3  # 工作节点ID位数（原5位）
        self.datacenter_id_bits = 3  # 数据中心ID位数（原5位）
        self.sequence_bits = 6  # 序列号位数（原12位）

        # 计算最大值
        self.max_worker_id = -1 ^ (-1 << self.worker_id_bits)
        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_id_bits)

        # 移位偏移量
        self.worker_id_shift = self.sequence_bits
        self.datacenter_id_shift = self.sequence_bits + self.worker_id_bits
        self.timestamp_left_shift = self.sequence_bits + self.worker_id_bits + self.datacenter_id_bits
        self.sequence_mask = -1 ^ (-1 << self.sequence_bits)

        # 验证参数
        if worker_id > self.max_worker_id or worker_id < 0:
            raise ValueError(f"worker_id 必须在 0 和 {self.max_worker_id} 之间")
        if datacenter_id > self.max_datacenter_id or datacenter_id < 0:
            raise ValueError(f"datacenter_id 必须在 0 和 {self.max_datacenter_id} 之间")

        self.worker_id = worker_id
        self.datacenter_id = datacenter_id
        self.sequence = 0
        self.last_timestamp = -1
        self.lock = threading.Lock()

    def _current_time(self):
        """获取当前时间戳（毫秒）"""
        return int(time.time() * 1000)

    def _wait_next_millis(self, last_timestamp):
        """等待下一毫秒"""
        timestamp = self._current_time()
        while timestamp <= last_timestamp:
            timestamp = self._current_time()
        return timestamp

    def next_id(self):
        """生成14位雪花ID"""
        with self.lock:
            timestamp = self._current_time()

            if timestamp < self.last_timestamp:
                raise Exception(f"时钟回退。拒绝生成ID达 {self.last_timestamp - timestamp} 毫秒")

            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.sequence_mask
                if self.sequence == 0:
                    timestamp = self._wait_next_millis(self.last_timestamp)
            else:
                self.sequence = 0

            self.last_timestamp = timestamp

            # 组合各部分生成最终ID
            return ((timestamp - 1742620000000) << self.timestamp_left_shift) | \
                (self.datacenter_id << self.datacenter_id_shift) | \
                (self.worker_id << self.worker_id_shift) | \
                self.sequence


if __name__ == "__main__":
    generator = SnowflakeIDGenerator(datacenter_id=1, worker_id=1)
    for _ in range(100):
        snowflake_id = generator.next_id()
        print(f"生成的14位雪花ID: {snowflake_id} 位数：{len(str((snowflake_id)))}")
